import fuzzywuzzy.fuzz as fz
import numpy as np
from numpy.core.defchararray import count
import pandas as pd
import os


def word_match_test(s1, s2) -> bool:
    if len(s2) < len(s1):
        s = s2
        s2 = s1
        s1 = s
    for c in s1:
        if c not in s2:
            return False
    return True


# build agent-alias rules based on worker-agents data
# assume that a worker works only for one agent
# may use ailias_read_from_manual() to extrate bad ones

# note it's maybe not right in case of [w1:a1,a2],[w2:a2,a1] .etc
def alias_build_rule(worker_to_agent: dict, worker_count: pd.DataFrame, file_name='banks/agent_alias_rule.txt',
                     f2_name='banks/agent_alias_false_rule.txt') -> list:
    f = open(file_name, 'w')
    f1 = open(f2_name, 'w')
    agent_alias_book = {}
    alias_agent_book = {}

    for work in worker_count.index:  # worker_to_agent:
        if work not in worker_to_agent:
            continue
        agents = worker_to_agent[work][0]
        if len(agents) < 2:
            continue
        # many names, maybe alias
        alias = []
        mainAgent = agents[0]
        if mainAgent[-1] == '院':  # 去除院字匹配
            mainAgent = mainAgent[:-1]
        if mainAgent[-1] == '学' and len(mainAgent) > 2:  # xx学院
            mainAgent = mainAgent[:-1]

        ratio_threshold = 50
        # if len(mainAgent) < 3:
        #     ratio_threshold = 50
        for ag in agents[1:]:
            d = fz.partial_ratio(mainAgent, ag)
            if d > ratio_threshold or word_match_test(mainAgent, ag):
                alias.append(ag)
                # alias_agent_book[ag] = agents[0]
            else:
                # print(d, agents[0],'\t',ag,'\t')
                f1.write(str(d) + ',' + agents[0] + '\t' + ag + '\n')
        if alias:
            for ag in alias:
                if agents[0] in alias_agent_book:  # that means this is a alias-formal set
                    major = alias_agent_book[agents[0]]
                    # assert major in agent_alias_book
                    if ag != major and ag not in agent_alias_book[major]:
                        # new set comes
                        agent_alias_book[major].append(ag)
                        alias_agent_book[ag] = major
                else:  # new set
                    if agents[0] not in agent_alias_book:
                        agent_alias_book[agents[0]] = [ag]
                        alias_agent_book[ag] = agents[0]
                    elif ag not in alias_agent_book:
                        agent_alias_book[agents[0]].append(ag)
                        alias_agent_book[ag] = agents[0]

            # if agents[0] not in agent_alias_book:
            #     agent_alias_book[agents[0]] = alias
            # else:
            #     agent_alias_book[agents[0]].extend(alias)

    for agent in agent_alias_book:
        alias = set(agent_alias_book[agent])
        n = len(alias)
        f.write(agent + ':')
        for i, a in enumerate(alias):
            f.write(a)
            if i < n - 1:
                f.write(',')
        f.write('\n')
    f.close()
    f1.close()
    return agent_alias_book, alias_agent_book


# include/exclude rule from file


def ailias_read_from_manual(agent_alias_book: dict, alias_agent_book: dict, f='banks/agent_alias_manual.txt'):
    if not os.path.exists(f):
        open(f, 'x').close()
        return
    rules = open(f).read().replace('，', ',')
    rules.replace('：', ':')
    rules.replace(' ', '')
    rules = rules.splitlines()

    mode = 'include'
    for rule in rules:
        if rule == '':
            continue
        if rule[0] == '#':
            mode = 'exclude' if 'exclude' in rule else 'include'
            continue
        if ':' not in rule:
            continue
        ind = rule.index(':')
        mainAgent = rule[:ind]
        items = rule[ind + 1:].split(',')
        if mode == 'include':
            if mainAgent in agent_alias_book:
                # need to set(list), but no matter
                for item in items:
                    if item not in agent_alias_book[mainAgent]:
                        agent_alias_book[mainAgent].append(item)

            else:
                agent_alias_book[mainAgent] = (items)
            for item in items:
                alias_agent_book[item] = mainAgent
        else:
            for item in items:
                if item in alias_agent_book:
                    alias_agent_book.pop(item)

            if mainAgent in agent_alias_book:  # exclude
                original = agent_alias_book[mainAgent]
                for i in items:
                    if i not in original:
                        continue
                    original.remove(i)
                if len(original) == 0:
                    agent_alias_book.pop(mainAgent)

    # return agent_alias_book, alias_agent_book


# read agent-alias rules from file
# use ailias_read_from_manual() to extrate bad ones
def alias_read_from(f='banks/agent_alias_rule.txt') -> dict:
    agent_alias_book = {}
    alias_agent_book = {}
    rules = open(f).read().replace('，', ',')
    rules.replace('：', ':')
    rules.replace(' ', '')
    rules = rules.splitlines()

    for rule in rules:
        ind = rule.index(':')
        mainAgent = rule[:ind]
        items = rule[ind + 1:].split(',')
        if mainAgent in agent_alias_book:
            # need to set(list), but no matter
            agent_alias_book[mainAgent].extend(items)
        else:
            agent_alias_book[mainAgent] = (items)
        for item in items:
            alias_agent_book[item] = mainAgent
    return agent_alias_book, alias_agent_book


# %% build agent-count map
# considering that a agent has many alias
def banks_load_for(name):
    fs = {'agent_alias': './banks/agent_alias_book.npy',
          'alias_agent': './banks/alias_agent_book.npy',
          'agent_worker': './banks/agent_to_worker.npy',
          'agent_count': './banks/agent_count.csv',
          'all_agent_count':'./banks/agents_count.csv'}
    if name in fs:
        dir = fs[name]
    else:
        dir = 'banks/bank_' + name + '.csv'
    if not os.path.exists(dir):
        return None
    if dir.endswith('csv'):
        return pd.read_csv(dir, index_col=0)
    if dir.endswith('.npy'):
        return np.load(dir, allow_pickle=True).item()
    return None


def statistics_agent_count(agents: dict, alias: dict, file='banks/agent_count.csv') -> pd.Series:
    counts = {}
    for agent in agents:
        c = agents[agent]  # count
        if agent in alias:
            formal = alias[agent]
        else:
            formal = agent

        if formal in counts:
            counts[formal] += c
        else:
            counts[formal] = c
    counts = pd.Series(counts, name='count')
    counts.to_csv(file, header=True)
    return counts
